package com.traitswu.websocket.dmxlog.core;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.websocket.Session;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * Created by traits on 2017/5/29.
 */


public class DmxLogWebSocketPipe {

    private static final Logger LOGGER = LoggerFactory.getLogger(DmxLogWebSocketPipe.class);

    private Map<String, Set<Session>> sessionMap = new HashMap<>();

    private ReentrantReadWriteLock readWriteLock = new ReentrantReadWriteLock();
    private Lock readeLock = readWriteLock.readLock();
    private Lock writeLock = readWriteLock.writeLock();


    private DmxLogWebSocketPipe() {
    }


    public void register(String loggerName, Session session) {
        executeWithLock(writeLock, () -> {
            Set<Session> sessions = sessionMap.computeIfAbsent(loggerName, key -> new HashSet<>());
            sessions.add(session);
        });
    }

    public void cancel(String loggerName, Session session) {
        executeWithLock(writeLock, () -> {
            Set<Session> sessions = sessionMap.get(loggerName);
            if (sessions == null) {
                return;
            }
            sessions.remove(session);
            if (sessions.isEmpty()) {
                sessionMap.remove(loggerName);
            }
        });

    }

    public void cancel(Session session) {
        executeWithLock(writeLock, () -> {
            Iterator<Map.Entry<String, Set<Session>>> it = sessionMap.entrySet().iterator();
            for (; it.hasNext(); ) {
                Set<Session> sessions = it.next().getValue();
                sessions.remove(session);
                if (sessions.isEmpty()) {
                    it.remove();
                }
            }
        });

    }


    public void write(String loggerName, byte[] message) {
        executeWithLock(readeLock, () -> {
            if (sessionMap.containsKey(loggerName)) {
                sessionMap.get(loggerName).forEach(session -> {
                    try {
                        //if client not close correctly, the session may be invalid
                        if (session.isOpen()) {
                            session.getBasicRemote().sendText(new String(message));
//                            session.getBasicRemote().sendBinary(ByteBuffer.wrap(message));
                        }
                    } catch (Exception e) {
                        if (!session.isOpen()) {
                            LOGGER.info("session had been closed.");
                        }
                        LOGGER.error("Unkown error: " + e.getMessage());
                    }
                });
            }

        });
    }


    public void read() {

    }


    public void clearCloseSession() {
        executeWithLock(writeLock, () -> {
            Iterator<Map.Entry<String, Set<Session>>> it = sessionMap.entrySet().iterator();
            for (; it.hasNext(); ) {
                Map.Entry<String, Set<Session>> entry = it.next();
                Set<Session> sessions = Optional.ofNullable(entry.getValue()).orElse(new HashSet<>());
                Iterator<Session> sessionIterator = sessions.iterator();
                for (; sessionIterator.hasNext(); ) {
                    Session session = sessionIterator.next();
                    if (!session.isOpen()) {
                        sessionIterator.remove();
                    }
                }
                if (sessions.isEmpty()) {
                    it.remove();
                }
            }
        });
    }


    private void executeWithLock(Lock lock, Runnable func) {
        try {
            lock.lock();
            func.run();
        } finally {
            lock.unlock();
        }
    }


    private static DmxLogWebSocketPipe pipe = new DmxLogWebSocketPipe();

    public static DmxLogWebSocketPipe getInstance() {
        return pipe;
    }

}
